import time
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import precision_score, recall_score, f1_score

import joblib
import pandas as pd

from TimeUtils import get_formatted_time

def train():
    # 记录训练开始时间
    start_time = time.time()

    #读取总的csv文件
    all_train_data = pd.read_csv('./all_train_data.csv')

    # 编码时间信息 秒为单位
    all_train_data['time'] = pd.to_datetime(all_train_data['time'])
    all_train_data['time'] = all_train_data['time'].astype('int64')

    # 数据预处理和特征工程
    # 这里需要根据您的数据集进行特征选择和处理，包括标签编码、缺失值处理、特征缩放等
    sampled_data = all_train_data.sample(frac=0.08, random_state=42)

    # 划分特征和标签
    X = sampled_data.drop(['渔船ID','type'], axis=1)
    y = sampled_data['type']


    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

    # 创建随机森林分类器
    rf_classifier = RandomForestClassifier(n_estimators=100, random_state=42)

    # 训练模型
    rf_classifier.fit(X_train, y_train)

    # 预测
    y_pred = rf_classifier.predict(X_test)

    # 评估模型性能
    precision = precision_score(y_test, y_pred, average='weighted')
    recall = recall_score(y_test, y_pred, average='weighted')
    f1 = f1_score(y_test, y_pred, average='weighted')

    print(f'Precision: {precision}')
    print(f'Recall: {recall}')
    print(f'F1 Score: {f1}')


    # 记录训练结束时间
    end_time = time.time()

    # 计算训练时间
    
    print(f"Training took {get_formatted_time(start_time, end_time)} seconds.")

    #保存模型
    joblib.dump(rf_classifier, 'trained_model-0.08.pkl')

if __name__ == '__main__':
    train()